//
// Ephi - simulation of magnetic fields and particles
// Copyright (C) 2007 Indrek Mandre <indrek(at)mare.ee>
// For more information please see http://www.mare.ee/indrek/ephi/
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License along
// with this program; if not, write to the Free Software Foundation, Inc.,
// 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
//

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <assert.h>
#include <vector>

#include "ephi.hpp"

#define RADIUS 0.15
#define WR 0.035
#define SPACING 0.01
#define COILPOT 10000
#define CURRENT 200000
#define COIL_MEMORY_USAGE 1000
#define CAGE_RADIUS 0.37

#if 1

#define PFIELD_RESOLUTION 0.004
#define PMAP_RESOLUTION   0.006
#define OCTREE_RESOLUTION 0.004
#define PMAP_SMOOTH 0.01

#else

#define PFIELD_RESOLUTION 0.01
#define PMAP_RESOLUTION 0.01
#define OCTREE_RESOLUTION 0.009
#define PMAP_SMOOTH 0.02
#endif

#define STAGES 100

#define CHARGE 8e-8

#define PARTICLES 10000
#define INJECTSPREAD 0.005
#define SPREADCOUNT 20000
#define SAMPLECOUNT 11
#define SAMPLEINTERVAL 2000

struct p_t
{
  vect3d pos, v;
};

struct stage_t
{
  enum { PCOUNT = PARTICLES * SAMPLECOUNT };
  stage_t () { particles.resize (PCOUNT); }
  std::vector<p_t> particles;
};
typedef std::list<stage_t> stages_t;
stages_t stages;

class Wb1OctreeBuilder : public OctreeBuilder
{
public:
  Wb1OctreeBuilder (const vect3d& c0, const vect3d& c1, prec_t resolution)
    : OctreeBuilder (c0, c1, 1, 2, 100, resolution), pmap(0)
  {
  }

  void reset (PMap& pmap) { this->pmap = &pmap; }

  void addValues (const vect3d& pos, CubicScalar *scalars, CubicVector *vectors)
  {
    pmap->addCubicValues (pos, scalars[0], vectors[0], vectors[1]);
  }

  bool inContact (const vect3d& pos)
  {
    return false;
  }

private:
  PMap *pmap;
};

int main (int argc, char *argv[])
{
  Ephi::setDebugLevel (Ephi::DEBUG_APP);

  // set up the bounding planes
  std::vector<vect3d> planes;
  planes.push_back (vect3d(0, 1, 0));
  planes.push_back (vect3d(1, 0, -1).normal());
  planes.push_back (vect3d(0, -1, 1).normal());

  PField pfield(PField::POINTSYSTEM);

  prec_t rr = RADIUS + (2 * WR + SPACING) / PREC_SQRT2;
  pfield.setResolution (PFIELD_RESOLUTION);
  pfield.addTorus (vect3d(-rr, 0, 0), vect3d(1, 0, 0), RADIUS, WR, COILPOT);
  pfield.addTorus (vect3d(rr, 0, 0), vect3d(-1, 0, 0), RADIUS, WR, COILPOT);
  pfield.addTorus (vect3d(0, -rr, 0), vect3d(0, 1, 0), RADIUS, WR, COILPOT);
  pfield.addTorus (vect3d(0, rr, 0), vect3d(0, -1, 0), RADIUS, WR, COILPOT);
  pfield.addTorus (vect3d(0, 0, rr), vect3d(0, 0, -1), RADIUS, WR, COILPOT);
  pfield.addTorus (vect3d(0, 0, -rr), vect3d(0, 0, 1), RADIUS, WR, COILPOT);

  // and the surrounding cage that compresses the potential gradient
  prec_t ores = pfield.setResolution (CAGE_RADIUS / 20.0);
  vect3d n1 = (planes[1] % planes[0]).normal();
  vect3d n2 = (planes[0] % planes[2]).normal();
  vect3d n3 = (planes[2] % planes[1]).normal();
  pfield.addSphereSegment (CAGE_RADIUS, n1, n2, n3, 0, true, 2);
  pfield.setResolution(ores);
  pfield.set_no_cdist (false);

  // get the 48-reduction transformations
  std::vector<transf3d> symm;
  std::vector<transf3d> rsymm;
  {
      std::vector<transf3d> transforms;
      transforms.reserve (48);
      const transf3d_pair *tprs = get_cube_polywell_transforms48 ();
      for ( size_t i = 0; i < 48; i++ )
        {
          if ( !(i & 1) )
              rsymm.push_back (tprs[i].f);
          else
              symm.push_back (tprs[i].f);
          transforms.push_back (tprs[i].f);
        }
      pfield.setTransforms (planes, transforms);
  }

  // make the octree builder
  vect3d octree_c0 = vect3d(0, 0, 0);
  vect3d octree_c1 = vect3d(CAGE_RADIUS + 0.02, CAGE_RADIUS + 0.02, CAGE_RADIUS + 0.02);
  Wb1OctreeBuilder builder(octree_c0, octree_c1, OCTREE_RESOLUTION);
  builder.setLimitRadius (CAGE_RADIUS + 0.01);
  for ( size_t i = 0; i < planes.size(); i++ )
      builder.addLimit (planes[i]);

  prec_t bmul = 1e6, emul = 1, qmul = 1;

  qmul = CHARGE / (STAGES * stage_t::PCOUNT * 48);

  Octree tree;
  Statics statics;
  // load or generate the bicubic coil field interpolator
  CoilData cd;
  if ( !cd.load ("wb1.cgen") )
    {
      {
        Ephi::debug (Ephi::DEBUG_APP, "Generating the coil bicubic interpolator and saving into wb1.cgen\n");
        CoilGen gen;
        gen.radius = RADIUS;
        gen.wr = WR;
        gen.memory_usage = COIL_MEMORY_USAGE;
        gen.do_efield = false;
        gen.do_cubic = true;
        gen.coil_matrix_n = 1;
        gen.set_space (SPACING);
        gen.geom_radial = (CAGE_RADIUS + 0.02) * 1.42;
        gen.geom_axial = (CAGE_RADIUS + 0.02) * 1.42;
        gen.generate ("wb1.cgen");
      }

      if ( !cd.load ("wb1.cgen") )
        {
          Ephi::debug (Ephi::DEBUG_CRITICAL, "%s: unable to read in the generated wb1.cgen!\n", argv[0]);
          exit(-1);
        }
    }

  make_polywell_cube (statics, RADIUS, WR, SPACING, CURRENT, 0, CoilDataFactory<StaticCICU>(cd));
  statics.addStaticElement (new StaticOctree48 (tree, true, true, true, bmul, emul));

  vect3d pmap_box(CAGE_RADIUS + 0.01, CAGE_RADIUS + 0.01, CAGE_RADIUS + 0.01);

  size_t N = 1;

  while (1)
    {
      char buf[128];
      sprintf (buf, "wb1_%06d_data.bin", (int)N);
      FILE *fp = fopen (buf, "rb");
      if ( !fp )
          break;
      FPReader r(fp);
      size_t c;
      r.read (c);
      if ( c != stage_t::PCOUNT )
          break;
      stages.push_back (stage_t());
      if ( stages.size() > STAGES )
          stages.pop_front();
      r.read (&stages.back().particles[0], sizeof (p_t) * stage_t::PCOUNT);
      printf ("Loaded ready-made stage %d\n", (int)N);
      N++;
    }

  while (1)
    {
      PMap pmap(pmap_box, PMAP_RESOLUTION, 0, PMAP_SMOOTH);

      for ( stages_t::iterator it = stages.begin(); it != stages.end(); it++ )
        {
          for ( size_t i = 0; i < stage_t::PCOUNT; i++ )
            {
#if 0
              // XXX no bfield for now
              for ( size_t j = 0; j < symm.size(); j++ )
                  pmap.addQ (symm[j] * it->particles[i].pos, vect3d(0, 0, 0), -qmul);
              for ( size_t j = 0; j < rsymm.size(); j++ )
                  pmap.addQ (rsymm[j] * it->particles[i].pos, vect3d(0, 0, 0), -qmul);
#else
              for ( size_t j = 0; j < symm.size(); j++ )
                  pmap.addQ (symm[j] * it->particles[i].pos, symm[j] * it->particles[i].v, -qmul);
              for ( size_t j = 0; j < rsymm.size(); j++ )
                  pmap.addQ (rsymm[j] * it->particles[i].pos, -(rsymm[j] * it->particles[i].v), -qmul);
#endif
            }
        }

      {
        PMap potmap(pmap_box, pmap_box.x / 10.0);
        for ( stages_t::iterator it = stages.begin(); it != stages.end(); it++ )
          {
            for ( size_t i = 0; i < stage_t::PCOUNT; i++ )
              {
                for ( size_t j = 0; j < symm.size(); j++ )
                    potmap.addQ (symm[j] * it->particles[i].pos, symm[j] * it->particles[i].v, -qmul);
                for ( size_t j = 0; j < rsymm.size(); j++ )
                    potmap.addQ (rsymm[j] * it->particles[i].pos, -(rsymm[j] * it->particles[i].v), -qmul);
              }
          }
        potmap.recalc();
        Statics st;
        st.addStaticElement (new StaticPMapCrudeEfield (potmap));
        pfield.setStatics (st);
        pfield.equalize ();
        pfield.fill (pmap);
      }

      pmap.recalc();
      builder.reset (pmap);
      builder.rebuild (tree);

      ElectroDynamics dyn(statics);
      dyn.setDisableFields (true);
      dyn.setDTM (0.01, 1);

      char buf[128];
      sprintf (buf, "wb1_%06d_", (int)N);
      std::string prefix = buf;
      Screen screen(800, 800, 2 * pmap_box.x, 2 * pmap_box.x);
      Scene scene(dyn);
      scene.calc (screen, Scene::CALC_BFIELD);
      scene.set_coloring (0, false);
      scene.render_map (screen);
      screen.write (prefix + "bfield.png");
      scene.set_coloring (32, true);
      scene.render_map (screen);
      screen.write (prefix + "bfield_log_stepped.png");
      scene.set_coloring (0, true);
      scene.render_lic (screen);
      screen.write (prefix + "bfield_lic.png");
      scene.calc (screen, Scene::CALC_EFIELD);
      scene.set_coloring (0, false);
      scene.render_map (screen);
      screen.write (prefix + "efield.png");
      scene.set_coloring (32, true);
      scene.render_map (screen);
      screen.write (prefix + "efield_log_stepped.png");
      scene.set_coloring (0, true);
      scene.render_lic (screen);
      screen.write (prefix + "efield_lic.png");
      scene.calc (screen, Scene::CALC_POTENTIAL);
      scene.set_coloring (32, false);
      scene.render_map (screen);
      screen.write (prefix + "pot.png");

      screen.clear();
      statics.draw (screen);
      screen.draw_axes ();
      for ( stages_t::iterator it = stages.begin(); it != stages.end(); it++ )
        {
          for ( size_t i = 0; i < stage_t::PCOUNT; i++ )
            {
              for ( size_t j = 0; j < symm.size(); j++ )
                  screen.set_pixel (symm[j] * it->particles[i].pos, Screen::RED);
              for ( size_t j = 0; j < rsymm.size(); j++ )
                  screen.set_pixel (rsymm[j] * it->particles[i].pos, Screen::RED);
            }
        }
      screen.write (prefix + "particles.png");

      Screen screen2(50, 50, 1.05 * 2 * pmap_box.x, 1.05 * 2 * pmap_box.x);
      scene.calc (screen2, Scene::CALC_POTENTIAL);
      scene.set_coloring (32, false);
      scene.render_map (screen2);
      scene.renderGnuplotData (screen2, (prefix + "gp_pot.bin").c_str());

      // inject the particles
      size_t injected = 0;
      while ( injected < PARTICLES )
        {
          for ( size_t i = 0; i < PARTICLES / 1000 + (PARTICLES < 1000 ? 1 : 0); i++ )
            {
              vect3d pos(0.25, 0, 0);
              pos.x += prec_t_drand() * INJECTSPREAD - INJECTSPREAD / 2;
              pos.y += prec_t_drand() * INJECTSPREAD - INJECTSPREAD / 2;
              pos.z += prec_t_drand() * INJECTSPREAD - INJECTSPREAD / 2;
              prec_t pot = dyn.getPotential (pos);
              pot -= 1000;
              dyn.inject_electron (pos, vect3d(-1, 0, 0), pot);
              injected++;
              if ( injected == PARTICLES )
                  break;
            }
          dyn.step();
          printf ("inject step: %d\n", (int)dyn.size());
        }

      dyn.step(SPREADCOUNT);
      printf ("spread done\n");

      size_t sindex = 0;

      stages.push_back (stage_t());
      if ( stages.size() > STAGES )
          stages.pop_front();
      for ( size_t i = 0; i < dyn.size(); i++ )
        {
          stages.back().particles[sindex].pos = dyn.get_pos(i);
          stages.back().particles[sindex].v = dyn.get_v(i);
          sindex++;
        }

      for ( size_t j = 1; j < SAMPLECOUNT; j++ )
        {
          dyn.step (SAMPLEINTERVAL);
          for ( size_t i = 0; i < dyn.size(); i++ )
            {
              stages.back().particles[sindex].pos = dyn.get_pos(i);
              stages.back().particles[sindex].v = dyn.get_v(i);
              sindex++;
            }
          printf ("sampleinterval stage %d done\n", (int)j);
        }

      FILE *fp = fopen ((prefix + "data.bin").c_str(), "wb");
      if ( fp )
        {
          FPWriter w(fp);
          w.write ((size_t)stage_t::PCOUNT);
          w.write (&stages.back().particles[0], sizeof (p_t) * stage_t::PCOUNT);
        }
      N++;
    }

  return 0;
}

